{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rStar-Math Integration Examples\n",
    "\n",
    "This notebook demonstrates how to integrate rStar-Math with various frameworks and platforms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "import os\n",
    "from typing import Dict, List\n",
    "from src.core.mcts import MCTS\n",
    "from src.core.ppm import ProcessPreferenceModel\n",
    "from src.models.model_interface import ModelFactory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Basic Integration\n",
    "\n",
    "First, let's see how to use rStar-Math directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# Initialize components\n",
    "mcts = MCTS.from_config_file('config/default.json')\n",
    "ppm = ProcessPreferenceModel.from_config_file('config/default.json')\n",
    "model = ModelFactory.create_model(\n",
    "    'openai',\n",
    "    os.getenv('OPENAI_API_KEY'),\n",
    "    'config/default.json'\n",
    ")\n",
    "\n",
    "# Solve a problem\n",
    "problem = \"What is the derivative of f(x) = x^2 + 3x?\"\n",
    "action, trajectory = mcts.search(problem)\n",
    "\n",
    "# Print solution steps with confidence scores\n",
    "for step in trajectory:\n",
    "    confidence = ppm.evaluate_step(step['state'], model)\n",
    "    print(f\"Step: {step['state']}\")\n",
    "    print(f\"Confidence: {confidence:.2f}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. LangChain Integration\n",
    "\n",
    "Here's how to use rStar-Math with LangChain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "from langchain.llms import OpenAI\n",
    "from examples.langchain_integration import RStarMathChain\n",
    "\n",
    "# Initialize LangChain components\n",
    "llm = OpenAI(temperature=0.7)\n",
    "chain = RStarMathChain(\n",
    "    llm=llm,\n",
    "    api_key=os.getenv('OPENAI_API_KEY')\n",
    ")\n",
    "\n",
    "# Solve problem\n",
    "result = chain.run(\"What is 2 + 2?\")\n",
    "\n",
    "print(\"Direct Solution:\")\n",
    "print(result['direct_solution'])\n",
    "print(f\"Score: {result['direct_score']:.2f}\\n\")\n",
    "\n",
    "print(\"Enhanced Solution:\")\n",
    "print(result['enhanced_solution'])\n",
    "print(f\"Score: {result['enhanced_score']:.2f}\\n\")\n",
    "\n",
    "print(f\"Improvement: {result['improvement']:.2%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Rasa Integration\n",
    "\n",
    "Example of using rStar-Math in a Rasa chatbot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "from rasa_sdk import Action\n",
    "from examples.rasa_integration import RStarMathAction\n",
    "\n",
    "# Create mock Rasa components for demonstration\n",
    "class MockDispatcher:\n",
    "    def utter_message(self, text: str):\n",
    "        print(f\"Bot: {text}\")\n",
    "\n",
    "class MockTracker:\n",
    "    def __init__(self, text: str):\n",
    "        self.latest_message = {\"text\": text}\n",
    "\n",
    "# Initialize action\n",
    "action = RStarMathAction()\n",
    "\n",
    "# Simulate conversation\n",
    "async def simulate_conversation():\n",
    "    dispatcher = MockDispatcher()\n",
    "    tracker = MockTracker(\"What is the derivative of x^2?\")\n",
    "    \n",
    "    print(f\"User: {tracker.latest_message['text']}\")\n",
    "    await action.run(dispatcher, tracker, {})\n",
    "\n",
    "# Run simulation\n",
    "import asyncio\n",
    "asyncio.run(simulate_conversation())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Custom Integration\n",
    "\n",
    "Example of creating a custom integration with rStar-Math."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "class MathTutor:\n",
    "    def __init__(self, api_key: str):\n",
    "        self.mcts = MCTS.from_config_file('config/default.json')\n",
    "        self.ppm = ProcessPreferenceModel.from_config_file('config/default.json')\n",
    "        self.model = ModelFactory.create_model('openai', api_key, 'config/default.json')\n",
    "        \n",
    "    def solve_with_explanation(self, problem: str) -> Dict:\n",
    "        # Get solution using rStar-Math\n",
    "        action, trajectory = self.mcts.search(problem)\n",
    "        \n",
    "        # Format steps with confidence scores\n",
    "        steps = []\n",
    "        total_confidence = 0.0\n",
    "        \n",
    "        for step in trajectory:\n",
    "            confidence = self.ppm.evaluate_step(step['state'], self.model)\n",
    "            steps.append({\n",
    "                'explanation': step['state'],\n",
    "                'confidence': confidence\n",
    "            })\n",
    "            total_confidence += confidence\n",
    "            \n",
    "        return {\n",
    "            'steps': steps,\n",
    "            'average_confidence': total_confidence / len(steps) if steps else 0.0\n",
    "        }\n",
    "\n",
    "# Test the custom integration\n",
    "tutor = MathTutor(os.getenv('OPENAI_API_KEY'))\n",
    "result = tutor.solve_with_explanation(\"Solve for x: 2x + 3 = 7\")\n",
    "\n",
    "print(\"Solution Steps:\")\n",
    "for i, step in enumerate(result['steps'], 1):\n",
    "    print(f\"\\nStep {i}:\")\n",
    "    print(f\"Explanation: {step['explanation']}\")\n",
    "    print(f\"Confidence: {step['confidence']:.2f}\")\n",
    "\n",
    "print(f\"\\nOverall Confidence: {result['average_confidence']:.2f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
